from torch import nn
import torch


# 定义一个类继承nn.Module
class MyNet(nn.Module):

    def __init__(self,num_emd=56,emd_dim=256,h_dim=128,lable_num=18):
        # num_emd：当前人名分类任务中对应的是去重后字符数量，经过统计是56
        # lable_num: 当前任务中对于的是国家种类数，经过统计是18
        # emb_dim: 软约束，对应的是词向量维度
        # h_dim: 软约束， 对应的是rnn隐藏层输出维度

        super().__init__() #初始化父类方法
        # 2.1 定义embedding层
        self.emb = nn.Embedding(num_embeddings=num_emd, # 注意向量的数量要大于索引最大值，否则超出边界
                                embedding_dim=emd_dim,
                                padding_idx=0,
                                max_norm=1)
        # 2.2 定义RNN层
        self.rnn = nn.GRU(input_size=emd_dim,
                          # 因为rnn层接收的是embedding的输出向量，
                          # 所以input_size=embedding层的输出向量维度也就是emd_dim
                          hidden_size=h_dim,# 此处设值无约束，是个超参
                          batch_first=True)
                            # 注意记得batch_first设成true
        # 2.3 定义输出的线性层
        self.l = nn.Linear(in_features=h_dim, # 因为这个线性层接收的rnn的输出向量， 所以in_features=rnn的输出向量维度，也就是rnn层的output_size
                           out_features=lable_num) # 此时不知道类别数量，它代表整个神经网络输出向量维度，暂时想设成多少就设成多少。

        # 声明一下规范化层
        self.ln = nn.LayerNorm(h_dim)

    def forward(self,batch_idx_tensor): #'人名' -> '逻辑向量'
        # batch_idx_tensor (batch,seq_len) [z, h, a, n, g] -> [1,2,3,4,5] -> [[0.1,0.2],[0.2,0.3]......]
        # 索引张量先经过embedding层得到批量序列张量
        batch_seq_tensor = self.emb(batch_idx_tensor) #( batch, seq_len, e_dim )
        # 之后将这个序列张量输入给rnn,只需记录hn即可
        # 但是注意rnn的输出有两个(output,hn)
        _, hn= self.rnn(batch_seq_tensor)
        hn = torch.squeeze(hn,dim=0) #dim=1代表在0轴上做挤压 # (batch,h_dim)
        hn = self.ln(hn)
        logits = self.l(hn)
        return logits # (batch,lable_num)


if __name__ == "__main__":
    # 测试用代码
    batch_idx_tensor = torch.LongTensor([[1,2,3],
                                         [2,3,4]]) #(batch=2,seq_len=3)

    net = MyNet()
    logits = net(batch_idx_tensor)
    print(logits.shape) #(batch=2,out_dim)